from collection.device_Info_collector import DeviceInfoCollector
from collection.device_dynamic_collector import DeviceDynamicCollector
from collection.nvidia_info_collector import NVIDIAGPUStaticCollector
from collection.nvidia_dynamic_collector import NVIDIAGPUDynamicCollector
from utils.config import Config
from utils.log_manager import LogManager
from utils.kafka_message_producer import KafkaMessageProducer
import time
import signal
import sys
import argparse
import uuid
from datetime import datetime


class DeviceCollector:
    def __init__(self, config_file: str):
        self.config = Config(config_file=config_file)
        self.log_manager = LogManager(self.config)
        self.kafka_producer = None
        self.topic = None
        self.signal_message_producer()
        self.device_info_collector = DeviceInfoCollector()
        self.device_dynamic_collector = DeviceDynamicCollector()
        self.nvidia_info_collector = NVIDIAGPUStaticCollector()
        self.nvidia_dynamic_collector = NVIDIAGPUDynamicCollector()
        self.running = True
        self.has_gpu = True
        self.collection_interval = self.config.get('collection_interval', 30)
        self.organization = self.config.get('organization')
        self.baseboard_serial_number = None

        # 注册信号处理
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)

    def signal_message_producer(self):
        kafka_host = self.config.get('kafka_host')
        if kafka_host:
            print("加载kafka消息生产者")
            username = self.config.get('kafka_username')
            password = self.config.get('kafka_password')
            """设置消息生产者"""
            if username and password:
                self.kafka_producer = KafkaMessageProducer(
                    bootstrap_servers=kafka_host,
                    sasl_plain_username = username,
                    sasl_plain_password = password
                )
                # if not self.kafka_producer.validate_kafka_auth():
                #     print("kafka认证失败,退出")
                #     sys.exit(1)
            else:
                self.kafka_producer = KafkaMessageProducer(
                    bootstrap_servers=kafka_host
                )
            self.topic = self.config.get('topic')

    @staticmethod
    def success_callback(record_metadata):
        print(f"异步发送成功: Topic={record_metadata.topic}, "
              f"Partition={record_metadata.partition}, "
              f"Offset={record_metadata.offset}")

    def send_message(self, key, message):
        self.kafka_producer.send_async(self.topic, message, key=key, callback=self.success_callback)

    def signal_handler(self, signum, frame):
        """处理退出信号"""
        print(f"\n收到退出信号 {signum}，正在安全停止...")
        self.running = False
        self.kafka_producer.close()

    def collect_static_info(self):
        """采集并记录静态信息"""
        print("开始采集设备静态信息...")
        try:
            self.device_info_collector.collect()
            device_info = self.device_info_collector.p_info()
            device_info['organization'] = self.organization
            # 采集gpu静态信息
            gpu_info = self.nvidia_info_collector.collect_static_info()
            if gpu_info['status'] == 'no_driver':
                self.has_gpu = False
            device_info['gpu_info'] = gpu_info
            self.log_manager.write_static_info(device_info)
            if self.kafka_producer:
                self.send_message(key="device_info", message=device_info)
            return True
        except Exception as e:
            print(f"采集静态信息失败: {e}")
            return False

    def collect_dynamic_info(self, is_first_collection=False):
        """采集并记录动态信息"""
        try:
            if is_first_collection:
                print("初始化动态信息采集...")

            self.device_dynamic_collector.collect_all()
            gpu_dynamic_info = None
            if self.has_gpu:
                gpu_dynamic_info = self.nvidia_dynamic_collector.collect_dynamic_info()
            if not is_first_collection:
                # 写入完整动态信息
                device_state = self.device_dynamic_collector.p_info()
                time_based_uuid = uuid.uuid1()  # 唯一id
                device_state['id'] = str(time_based_uuid)
                device_state['organization'] = self.organization  # 组织
                device_state['baseboard_serial_number'] = self.log_manager.get_device_id()  # 设备唯一号
                if gpu_dynamic_info:
                    device_state['gpu_info'] = gpu_dynamic_info
                else:
                    device_state['gpu_info'] = {"status": "no_driver"}
                self.log_manager.write_dynamic_info(device_state, is_summary=False)
                if self.kafka_producer:
                    self.send_message(key="device_state", message=device_state)
                # # 写入摘要动态信息
                # device_state_summary = self.device_dynamic_collector.p_summary_info()
                # if gpu_dynamic_info:
                #     device_state_summary['gpu_info'] = gpu_dynamic_info
                # else:
                #     device_state_summary['gpu_info'] = {"status": "no_driver"}
                # self.log_manager.write_dynamic_info(device_state_summary, is_summary=True)

        except Exception as e:
            print(f"采集动态信息失败: {e}")

    def run(self):
        """主运行循环"""
        print("设备信息采集器启动")
        print(f"组织: {self.config.get('organization')}")
        print(f"采集间隔: {self.collection_interval}秒")
        print(f"日志目录: {self.config.get('log_directory')}")

        # 测试写入权限
        if not self.log_manager.test_write_access():
            print("错误: 无法写入日志目录，请检查权限")
            return

        # 检查今天是否需要采集静态信息
        if self.log_manager.need_static_info_today():
            print("今天还没有静态信息记录，开始采集...")
            if not self.collect_static_info():
                print("静态信息采集失败，程序退出")
                return
        else:
            print("今天已有静态信息记录")
            # device_info = self.device_info_collector.p_info()
            # self.baseboard_serial_number = device_info['baseboard_serial_number']
            # 采集gpu静态信息
            gpu_info = self.nvidia_info_collector.collect_static_info()
            if gpu_info['status'] == 'no_driver':
                self.has_gpu = False

        # 初始化动态采集器
        self.collect_dynamic_info(is_first_collection=True)

        print(f"等待 {self.collection_interval} 秒后开始定期采集...")
        time.sleep(self.collection_interval)

        last_check_date = datetime.now().date()
        collection_count = 0

        while self.running:
            try:
                current_date = datetime.now().date()

                # 检查是否是新的一天
                if current_date != last_check_date:
                    print(f"新的一天开始: {current_date}")
                    if self.log_manager.need_static_info_today():
                        self.collect_static_info()
                    last_check_date = current_date
                    collection_count = 0

                # 采集动态信息
                collection_count += 1
                current_time = datetime.now().strftime('%H:%M:%S')
                print(f"[{current_time}] 第 {collection_count} 次采集...")

                self.collect_dynamic_info()

                # 等待下次采集
                time.sleep(self.collection_interval)

            except KeyboardInterrupt:
                print("\n用户中断程序")
                break
            except Exception as e:
                print(f"采集过程中发生错误: {e}")
                time.sleep(10)

        print("设备信息采集器已停止")

        # 显示最终日志文件信息
        current_log_file = self.log_manager.get_current_log_file()
        if current_log_file and current_log_file.exists():
            file_size = current_log_file.stat().st_size
            print(f"日志文件: {current_log_file.name} (大小: {file_size} 字节)")

def create_parser():
    parser = argparse.ArgumentParser(
        description='设备信息采集工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
    使用示例:
      %(prog)s --config-file config.json             # 配置文件路径
            """
    )

    parser.add_argument(
        '--config-file', '-c',
        type=str,
        help='配置文件路径'
    )

    return parser


def main():
    parser = create_parser()
    args = parser.parse_args()
    config_file_path = args.config_file
    print(f"设置配置文件路径：{config_file_path}")
    try:
        collector = DeviceCollector(config_file=config_file_path if config_file_path else 'config.json')
        collector.run()
    except Exception as e:
        print(f"程序启动失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
